#!/bin/bash


RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# 检查参数
if [ $# -eq 0 ]; then
    echo -e "${RED}错误：请指定地区代码参数，例如：./check_update_ipset.sh cn${NC}"
    exit 1
fi
# 定义 ipset 集合名称
ZONE="$1"
IPSET_NAME="${ZONE}_ips"


#==============================================================
# 检查 ipset 是否安装
check_ipset_installed() {
  if ! command -v ipset &> /dev/null; then
    echo -e "${RED}错误：ipset 未安装。请先安装 ipset (例如：sudo apt install ipset)${NC}"
    exit 1
  else
    echo -e "${GREEN}ipset 已安装。"
  fi
}

# 检查 iptables 是否安装
check_iptables_installed() {
  if ! command -v iptables &> /dev/null; then
    echo "错误：iptables 未安装。请先安装 iptables (例如：sudo apt install iptables)"
    exit 1
  else
    echo -e "${GREEN}iptables 已安装。"
  fi
}

# 检查ufw是否安装,是否安装
check_ufw_installed() {
  if ! command -v ufw &> /dev/null; then
    echo "错误：ufw 未安装。请先安装 ufw (例如：sudo apt install ufw)"
    exit 1
  else
    echo -e "${GREEN}ufw 已安装。"
  fi
}
check_installed() {
  check_ipset_installed
  check_iptables_installed
  check_ufw_installed
}

#==================================================================
# 删除所有引用该 ipset 的 iptables 规则
delete_iptables_rules() {
  RULES=$(sudo iptables-save | grep "$IPSET_NAME")
  if [ -n "$RULES" ]; then
    echo "正在删除引用 $IPSET_NAME 的 iptables 规则..."
    while read -r rule; do
      chain=$(echo "$rule" | awk '{print $2}')
      params=$(echo "$rule" | sed -E "s/^-A [^ ]+ //")
      sudo iptables -D "$chain" $params
    done <<< "$RULES"
    else
      echo "没有找到引用 $IPSET_NAME 的 iptables 规则。"
  fi
}

# 销毁 ipset 集合
destroy_ipset() {
  echo -e "${YELLOW}正在销毁 ipset 集合 $IPSET_NAME ...${NC}"
  sudo ipset destroy "$IPSET_NAME"
  if [ $? -eq 0 ]; then
    echo -e "${GREEN}ipset 集合 $IPSET_NAME 已销毁。${NC}"
  else
    echo -e "${RED}销毁 ipset 集合 $IPSET_NAME 失败。${NC}"
  fi
}

# 创建并加载 ipset 集合及相关规则的函数
create_and_load_ipset() {
  local ipset_restore_file="zones/${IPSET_NAME}.ipset"
  local ufw_rules_file="/etc/ufw/before.rules"

  echo "创建 ipset 集合: $IPSET_NAME"
  if [ -f "$ipset_restore_file" ]; then
    echo "正在从 $ipset_restore_file 恢复 ipset 集合..."
    sudo ipset restore < "$ipset_restore_file"
  else
    echo -e "${RED}错误：$ipset_restore_file 文件不存在，请检查。${NC}"
    exit 1
  fi

  echo "清理旧的ufw规则..."
  sudo sed -i "/match-set $IPSET_NAME src/d" "$ufw_rules_file"
  
  echo "添加新的ufw规则..."
  sudo sed -i "/^COMMIT$/i -A ufw-before-input -m set --match-set $IPSET_NAME src -j ACCEPT" "$ufw_rules_file"
  
  # echo "重新加载 ufw 规则..."
  sudo ufw reload
}
# 检查 ipset 集合是否存在
check_and_update() {
  if ! sudo ipset list "$IPSET_NAME" &> /dev/null; then
    echo -e "${YELLOW}ipset 集合 $IPSET_NAME 不存在, 开始创建...${NC}"
  else
    echo -e "${GREEN}ipset 集合 $IPSET_NAME 存在, 开始更新...${NC}"
    delete_iptables_rules
    destroy_ipset
  fi
  create_and_load_ipset
}
# 检查 ufw-before-input 链是否存在
check_ufw_chain_exists() {
  if ! sudo iptables -L ufw-before-input &> /dev/null; then
    echo "警告：ufw-before-input 链不存在，可能 ufw 未启用或未正确配置。"
  else
    echo "ufw-before-input 链存在。"
    sudo iptables -L ufw-before-input -n --line-numbers
  fi
}
main() {
  check_installed
  check_and_update
  check_ufw_chain_exists
}

main 
exit 0


